package com.ygp.asyn.entity;

import com.ygp.asyn.util.FileUtils;
import lombok.extern.slf4j.Slf4j;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Map;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

@Slf4j
public class ZipFile {

    private ZipOutputStream zos;
    private ByteArrayOutputStream zipByOut;
    private ByteArrayInputStream zipByIn;

    public ZipFile() {
        this.zipByOut = new ByteArrayOutputStream();
        this.zos = new ZipOutputStream(this.zipByOut);
    }

    /**
     * 添加完成后需要，需要手动关闭流
     *
     * @param outputStream 被添加的流对象；
     * @param fileName     对应的文件名；文件名需要带后缀
     */
    public void add(ByteArrayOutputStream outputStream, String fileName) {
        try {
            ZipEntry zipEntry = new ZipEntry(fileName);
            zipEntry.setSize(outputStream.size());
            zipEntry.setCompressedSize(outputStream.size());
            zos.putNextEntry(zipEntry);
            outputStream.writeTo(zos);
        } catch (Exception e) {
            log.error("addToZipIn error:", e);
        } finally {
            FileUtils.closeStream(outputStream);
            FileUtils.setNull(outputStream);
            closeEntry(zos);
        }
    }

    /**
     * 获取最终结果
     *
     * @return ；
     */
    public ByteArrayInputStream getResultIn() {
        this.zipByIn = new ByteArrayInputStream(zipByOut.toByteArray());
        return this.zipByIn;
    }

    public byte[] getResultByte() {
        return zipByOut.toByteArray();
    }

    public ByteArrayOutputStream getResultOut() {
        return zipByOut;
    }

    /**
     * 添加完后 关闭流
     */
    public void close() {
        FileUtils.closeStream(this.zos, this.zipByOut, this.zipByIn);
    }

    public void setNull() {
        FileUtils.setNull(this.zos, this.zipByOut, this.zipByIn);
    }

    /**
     * 文件名需要带后缀
     *
     * @param map <文件名,流对象>
     * @return ；
     */
    public static InputStream byOutToZipIn(Map<String, ByteArrayOutputStream> map) {
        ByteArrayOutputStream zipByOut = new ByteArrayOutputStream();
        ZipOutputStream zos = new ZipOutputStream(zipByOut);
        map.forEach((key, value) -> {
            try {
                ZipEntry zipEntry = new ZipEntry(key);
                zipEntry.setSize(value.size());
                zipEntry.setCompressedSize(value.size());
                zos.putNextEntry(zipEntry);
                value.writeTo(zos);
            } catch (IOException e) {
                throw new RuntimeException(e);
            } finally {
                FileUtils.closeStream(value);
                closeEntry(zos);
            }
        });
        ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(zipByOut.toByteArray());
        FileUtils.closeStream(zos, zipByOut, byteArrayInputStream);
        return byteArrayInputStream;
    }


    public static void closeEntry(ZipOutputStream... zipOss) {
        for (ZipOutputStream zos : zipOss) {
            if (zos != null) {
                try {
                    zos.closeEntry();
                } catch (Exception e) {
                    log.error("closeEntry error", e);
                }
            }
        }
    }


}
